package com.hk.core.data.r2dbc.connection;

import org.reactivestreams.Publisher;
import org.springframework.r2dbc.connection.lookup.AbstractRoutingConnectionFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.Objects;

/**
 * R2dbc 多数据源
 *
 * @author Kevin
 * @date 2023-11-12 17:00
 */
public class DynamicRoutingConnectionFactory extends AbstractRoutingConnectionFactory {

    private static final String DB_KEY_NAME = "R_JDBC_DB_NAME";

    private static String defaultDbName;

    public DynamicRoutingConnectionFactory(String defaultDbName, Object defaultTargetConnectionFactory, Map<?, ?> targetConnectionFactories) {
        setDefaultTargetConnectionFactory(defaultTargetConnectionFactory);
        setTargetConnectionFactories(targetConnectionFactories);
        setDefaultDbName(defaultDbName);
    }

    private static void setDefaultDbName(String defaultDbName) {
        if (Objects.isNull(DynamicRoutingConnectionFactory.defaultDbName)) {
            DynamicRoutingConnectionFactory.defaultDbName = defaultDbName;
        }
    }

    public static <T> Publisher<T> putR2dbcDataSource(Publisher<T> publisher, Mono<String> dbName) {
        if (publisher instanceof Mono<T> mono) {
            return putMonoR2dbcDataSource(mono, dbName);
        } else if (publisher instanceof Flux<T> flux) {
            return putFluxR2dbcDataSource(flux, dbName);
        }
        return publisher;
    }

    public static <T> Mono<T> putMonoR2dbcDataSource(Mono<T> mono, Mono<String> dbName) {
        return dbName
                .defaultIfEmpty(defaultDbName)
                .flatMap(db -> mono.contextWrite(cxt -> cxt.put(DB_KEY_NAME, db)));
    }


    public static <T> Flux<T> putFluxR2dbcDataSource(Flux<T> flux, Mono<String> dbName) {
        return dbName
                .defaultIfEmpty(defaultDbName)
                .flatMapMany(db -> flux.contextWrite(cxt -> cxt.put(DB_KEY_NAME, db)));
    }

    @Override
    protected Mono<Object> determineCurrentLookupKey() {
        return Mono.deferContextual(Mono::just)
                .filter(cxt -> cxt.hasKey(DB_KEY_NAME))
                .map(cxt -> cxt.get(DB_KEY_NAME));
    }
}
